package segmentId;

import com.zaxxer.hikari.HikariDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @program: rocsun
 * @description: 号段模式分布式ID生成器， 支持双缓冲优化、动态步长调整、故障恢复等特性
 * @author: yx.guo
 * @create: 2025-08-10 19:28
 **/
public class SegmentIdGenerator {
    // 数据库配置
    private final DataSource dataSource;
    private final String bizTag;

    // 号段配置
    private volatile Segment currentSegment;
    private volatile Segment nextSegment;
    private volatile int step = 1000; // 默认步长

    // 状态控制
    private final Lock segmentLock = new ReentrantLock();
    private volatile boolean loadingNext = false;
    private final double loadThreshold = 0.2; // 20%时加载下一号段

    // 监控指标
    private final AtomicLong generatedIds = new AtomicLong(0);
    private final AtomicLong segmentSwitches = new AtomicLong(0);

    /**
     * 号段内部类
     */
    private static class Segment {
        final long start;
        final long end;
        final AtomicLong current;

        Segment(long start, long end) {
            this.start = start;
            this.end = end;
            this.current = new AtomicLong(start);
        }

        /**
         * 是否耗尽
         *
         * @return
         */
        boolean isExhausted() {
            return current.get() > end;
        }

        /**
         * 已使用比例
         *
         * @return
         */
        double usageRatio() {
            long used = current.get() - start;
            long total = end - start + 1;
            return (double) used / total;
        }
    }

    /**
     * 初始化ID生成器
     *
     * @param dataSource 数据库连接池
     * @param bizTag     业务标识（区分不同业务）
     */
    public SegmentIdGenerator(DataSource dataSource, String bizTag) {
        this.dataSource = dataSource;
        this.bizTag = bizTag;
        initializeSegments();
    }

    /**
     * 初始化号段
     */
    private void initializeSegments() {
        currentSegment = loadSegmentFromDB();
        nextSegment = loadSegmentFromDB();
        segmentSwitches.incrementAndGet();
    }

    /**
     * 从数据库加载号段（其实是加载新号段）
     */
    private Segment loadSegmentFromDB() {
        String sql = "UPDATE id_segment SET max_id = max_id + ? WHERE biz_tag = ?";
        String querySql = "SELECT max_id, step FROM id_segment WHERE biz_tag = ?";

        try (Connection conn = dataSource.getConnection();
            PreparedStatement updateStmt = conn.prepareStatement(sql);
            PreparedStatement queryStmt = conn.prepareStatement(querySql)) {

            // 开启事务
            conn.setAutoCommit(false);

            // 1. 更新数据库获取新号段
            updateStmt.setInt(1, step);
            updateStmt.setString(2, bizTag);
            int affectedRows = updateStmt.executeUpdate();

            if (affectedRows == 0) {
                // 记录不存在，初始化
                initializeSegmentInDB(conn);
                return loadSegmentFromDB();
            }

            // 2. 查询更新后的值
            queryStmt.setString(1, bizTag);
            try (ResultSet rs = queryStmt.executeQuery()) {
                if (rs.next()) {
                    long newMaxId = rs.getLong("max_id");
                    int dbStep = rs.getInt("step");

                    // 更新步长（如果数据库步长变化）
                    if (dbStep != step) {
                        step = dbStep;
                    }

                    long start = newMaxId - step + 1;
                    conn.commit();
                    return new Segment(start, newMaxId);
                }
            }

            conn.rollback();
            throw new RuntimeException("Failed to load segment for tag: " + bizTag);

        } catch (SQLException e) {
            throw new RuntimeException("Database error while loading segment", e);
        }
    }

    /**
     * 初始化数据库中的号段记录
     */
    private void initializeSegmentInDB(Connection conn) throws SQLException {
        String insertSql = "INSERT INTO id_segment (biz_tag, max_id, step, description) VALUES (?, ?, ?, ?)";

        try (PreparedStatement stmt = conn.prepareStatement(insertSql)) {
            stmt.setString(1, bizTag);
            stmt.setLong(2, step); // 初始值=步长
            stmt.setInt(3, step);
            stmt.setString(4, "Auto-initialized");
            stmt.executeUpdate();
            conn.commit();
        }
    }

    /**
     * 获取下一个ID
     */
    public long nextId() {
        // 尝试从当前号段获取ID，并加载下一号段
        long id = getIdFromCurrentSegment();
        // 如果存在，记录并返回
        if (id != -1) {
            generatedIds.incrementAndGet();
            return id;
        }

        // 当前号段用完，切换到下一号段，并且
        return switchSegmentAndGetId();
    }

    /**
     * 从当前号段获取ID，并检查加载下一号段
     *
     * @return ID 或 -1（号段耗尽）
     */
    private long getIdFromCurrentSegment() {
        Segment segment = currentSegment;
        if (segment == null) return -1;

        long id = segment.current.getAndIncrement();
        if (id <= segment.end) {
            // 检查是否需要预加载下一号段
            checkLoadNextSegment(segment);
            return id;
        }

        return -1;
    }

    /**
     * 检查是否需要加载下一号段
     */
    private void checkLoadNextSegment(Segment segment) {
        if (nextSegment == null &&
                !loadingNext &&
                segment.usageRatio() >= loadThreshold) {

            // 异步加载下一号段
            new Thread(() -> {
                try {
                    // 标记需要加载下一号段
                    loadingNext = true;
                    Segment newSegment = loadSegmentFromDB();

                    segmentLock.lock();
                    try {
                        nextSegment = newSegment;
                    } finally {
                        segmentLock.unlock();
                    }
                } finally {
                    loadingNext = false;
                }
            }).start();
        }
    }

    /**
     * 切换号段并获取ID
     */
    private long switchSegmentAndGetId() {
        segmentLock.lock();
        try {
            // 双重检查
            long id = getIdFromCurrentSegment();
            if (id != -1) return id;

            // 确保下一号段已加载
            if (nextSegment == null) {
                nextSegment = loadSegmentFromDB();
            }

            // 切换号段
            currentSegment = nextSegment;
            nextSegment = null;
            // 记录切换次数+1
            segmentSwitches.incrementAndGet();

            /*------------------从新号段获取id---------------------*/
            // 从新号段获取ID
            id = currentSegment.current.getAndIncrement();
            if (id > currentSegment.end) {
                throw new IllegalStateException("ID generation overflow for tag: " + bizTag);
            }

            // 异步加载下一个号段
            new Thread(this::loadNextSegmentAsync).start();
            // 记录一下使用了一次id
            generatedIds.incrementAndGet();
            return id;
        } finally {
            segmentLock.unlock();
        }
    }

    /**
     * 异步加载下一号段
     */
    private void loadNextSegmentAsync() {
        if (nextSegment == null && !loadingNext) {
            loadingNext = true;
            try {
                Segment newSegment = loadSegmentFromDB();

                segmentLock.lock();
                try {
                    nextSegment = newSegment;
                } finally {
                    segmentLock.unlock();
                }
            } finally {
                loadingNext = false;
            }
        }
    }

    /**
     * 调整步长（动态扩容）
     *
     * @param newStep 新步长
     */
    public void adjustStep(int newStep) {
        if (newStep <= 0) {
            throw new IllegalArgumentException("Step must be positive");
        }

        String sql = "UPDATE id_segment SET step = ? WHERE biz_tag = ?";

        try (Connection conn = dataSource.getConnection();
            PreparedStatement stmt = conn.prepareStatement(sql)) {

            stmt.setInt(1, newStep);
            stmt.setString(2, bizTag);
            int affected = stmt.executeUpdate();

            if (affected > 0) {
                this.step = newStep;
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to adjust step", e);
        }
    }

    /**
     * 获取监控指标
     */
    public String getMetrics() {
        return String.format("SegmentIdGenerator[%s] - Generated: %d, Segment Switches: %d, Current Step: %d",
                bizTag, generatedIds.get(), segmentSwitches.get(), step);
    }

    // 测试用例
    public static void main(String[] args) {
        // 初始化数据库连接池 (使用HikariCP示例)
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl("jdbc:mysql://117.72.78.44:3306/rocsun");
        dataSource.setUsername("root");
        dataSource.setPassword("123456");

        // 创建ID生成器（业务标签：order）
        SegmentIdGenerator idGenerator = new SegmentIdGenerator(dataSource, "order");

        // 生成10个ID
        for (int i = 0; i < 10; i++) {
            long id = idGenerator.nextId();
            System.out.println("生成的ID: " + id);
        }

        // 打印监控指标
        System.out.println(idGenerator.getMetrics());

        // 动态调整步长
        idGenerator.adjustStep(2000);
        System.out.println("调整步长为2000");

        // 再生成5个ID
        for (int i = 0; i < 5; i++) {
            long id = idGenerator.nextId();
            System.out.println("生成的ID: " + id);
        }

        // 打印最终监控指标
        System.out.println(idGenerator.getMetrics());
    }
}
